package algorithm.graph;

import algorithm.unionfind.general.UnionFind;
import datastructure.heap.BinaryHeap;

import java.util.*;


/**
 * @Description:
 * @Author Ammar
 * @Create 2023/5/19 18:43
 */
public class ListGraph<V, E> extends Graph<V, E> {

    private Map<V, Vertex<V, E>> vertices = new HashMap<>();
    private Set<Edge<V, E>> edges = new HashSet<>();
    private Comparator<Edge<V, E>> edgeComparator = (Edge<V, E> e1, Edge<V, E> e2) -> {
        return weightManager.compare(e1.weight, e2.weight);
    };

    public ListGraph() {
        this(null);
    }

    public ListGraph(WeightManager<E> weightManager) {
        super(weightManager);
    }


    @Override
    public int edgesSize() {
        return edges.size();
    }

    @Override
    public int verticesSize() {
        return vertices.size();
    }

    @Override
    public void addVertex(V v) {
        if (vertices.containsKey(v)) return;
        vertices.put(v, new Vertex<>(v));
    }

    @Override
    public void addEdge(V from, V to) {
        addEdge(from, to, null);
    }

    @Override
    public void addEdge(V from, V to, E weight) {
        Vertex<V, E> fromVertex = vertices.get(from);
        if (fromVertex == null) {
            fromVertex = new Vertex<>(from);
            vertices.put(from, fromVertex);
        }

        Vertex<V, E> toVertex = vertices.get(to);
        if (toVertex == null) {
            toVertex = new Vertex<>(to);
            vertices.put(to, toVertex);
        }

        Edge<V, E> edge = new Edge<>(fromVertex, toVertex);
        edge.weight = weight;

        if (fromVertex.outEdges.remove(edge)) {
            toVertex.inEdges.remove(edge);
            edges.remove(edge);
        }

        fromVertex.outEdges.add(edge);
        toVertex.inEdges.add(edge);
        edges.add(edge);
    }

    @Override
    public void removeVertex(V v) {
        Vertex<V, E> remove = vertices.remove(v);
        if (remove == null) return;
        Iterator<Edge<V, E>> outIterator = remove.outEdges.iterator();
        while (outIterator.hasNext()) {
            Edge<V, E> edge = outIterator.next();
            edge.to.inEdges.remove(edge);
            outIterator.remove();
            edges.remove(edge);
        }
        Iterator<Edge<V, E>> inIterator = remove.inEdges.iterator();
        while (inIterator.hasNext()) {
            Edge<V, E> edge = inIterator.next();
            edge.from.outEdges.remove(edge);
            inIterator.remove();
            edges.remove(edge);
        }
    }

    @Override
    public void removeEdge(V from, V to) {
        Edge<V, E> edge = new Edge<>(vertices.get(from), vertices.get(to));
        if (edges.remove(edge)) {
            vertices.get(from).outEdges.remove(edge);
            vertices.get(to).inEdges.remove(edge);
        }
    }

    @Override
    public void bfs(V begin, VertexVisitor<V> visitor) {
        Vertex<V, E> beginVertex = vertices.get(begin);
        if (beginVertex == null) return;

        Set<Vertex<V, E>> visitedVertices = new HashSet<>();
        Queue<Vertex<V, E>> queue = new LinkedList<>();
        queue.offer(beginVertex);
        visitedVertices.add(beginVertex);

        while (!queue.isEmpty()) {
            Vertex<V, E> vertex = queue.poll();
            if (visitor.visit(vertex.value)) return;

            for (Edge<V, E> edge : vertex.outEdges) {
                if (visitedVertices.contains(edge.to)) continue;
                queue.offer(edge.to);
                visitedVertices.add(edge.to);
            }
        }
    }

    // 递归
    @Override
    public void dfs(V begin, VertexVisitor<V> visitor) {
        Vertex<V, E> beginVertex = vertices.get(begin);
        if (beginVertex == null) return;
        dfs(beginVertex, new HashSet<>(), visitor);
    }

    private void dfs(Vertex<V, E> vertex, Set<Vertex<V, E>> visitedVertices, VertexVisitor<V> visitor) {
        if (visitor.visit(vertex.value)) return;
        visitedVertices.add(vertex);
        for (Edge<V, E> edge : vertex.outEdges) {
            if (visitedVertices.contains(edge.to)) continue;
            dfs(edge.to, visitedVertices, visitor);
        }
    }

    // 非递归
    @Override
    public void dfs1(V begin, VertexVisitor<V> visitor) {
        Vertex<V, E> beginVertex = vertices.get(begin);
        if (beginVertex == null) return;

        Set<Vertex<V, E>> visitedVertices = new HashSet<>();
        Stack<Vertex<V, E>> stack = new Stack<>();
        // 将起点入栈
        stack.push(beginVertex);
        visitedVertices.add(beginVertex);

        // 访问起点
        if (visitor.visit(beginVertex.value)) return;

        while (!stack.isEmpty()) {
            // 取出栈顶元素
            Vertex<V, E> vertex = stack.pop();
            // 遍历栈顶元素的出边
            for (Edge<V, E> edge : vertex.outEdges) {
                if (visitedVertices.contains(edge.to)) continue;

                stack.push(edge.from);
                stack.push(edge.to);
                visitedVertices.add(edge.to);
                if (visitor.visit(edge.to.value)) return;
                break;
            }
        }
    }

    // 拓扑排序
    @Override
    public List<V> topologicalSort() {
        List<V> list = new ArrayList<>(); // 存储最终拓扑排序的结果
        Queue<Vertex<V, E>> queue = new LinkedList<>(); // 存储入度为0的顶点
        Map<Vertex<V, E>, Integer> ins = new HashMap<>(); // 存储每个顶点的入度，随着遍历过程，入度会发生变化

        // 初始化（将度为0的顶点入队）
        vertices.forEach((V v, Vertex<V, E> vertex) -> {
            int in = vertex.inEdges.size();
            if (in == 0) {
                queue.offer(vertex);
            } else {
                ins.put(vertex, in);
            }
        });

        while (!queue.isEmpty()) {
            Vertex<V, E> vertex = queue.poll();
            list.add(vertex.value);

            for (Edge<V, E> edge : vertex.outEdges) {
                int toIn = ins.get(edge.to) - 1;
                if (toIn == 0) {
                    queue.offer(edge.to);
                } else {
                    ins.put(edge.to, toIn);
                }
            }
        }

        return list;
    }

    // 最小生成树
    @Override
    public Set<EdgeInfo<V, E>> mst() {
//        return prim();
        return kruskal();
    }

    // 最坏的复杂度 ElogE 稍好
    private Set<EdgeInfo<V, E>> prim() {
        // 随机选择一个顶点
        Iterator<Vertex<V, E>> it = vertices.values().iterator();
        // 如果没有顶点，返回 null
        if (!it.hasNext()) return null;
        // 存储最小生成树的边
        Set<EdgeInfo<V, E>> edgeInfos = new HashSet<>();
        // 存储已经访问过的顶点
        Set<Vertex<V, E>> addedVertices = new HashSet<>();
        Vertex<V, E> begin = vertices.values().iterator().next();
        addedVertices.add(begin);
        // 小顶堆，选最小值
        // 不使用官方的 PriorityQueue，在设置比较器时，无法批量建堆，会导致时间复杂度变为 O(nlogn)
        // 使用自己的
        BinaryHeap<Edge<V, E>> heap = new BinaryHeap<>(begin.outEdges, edgeComparator);
        int vertexSize = vertices.size();
        while (!heap.isEmpty() && addedVertices.size() < vertexSize) {
            // 拿出权值最小的边
            Edge<V, E> edge = heap.remove();
            // 如果这条边的 to 顶点已经访问过了，说明这条边无效，继续下一次循环
            if (addedVertices.contains(edge.to)) continue;
            edgeInfos.add(edge.info());
            addedVertices.add(edge.to);
            // 将to顶点的 outEdges 添加到heap中，此时会把 edge 又添加进去
            heap.addAll(edge.to.outEdges);
        }

        return edgeInfos;
    }

    // ElogE 复杂度
    private Set<EdgeInfo<V, E>> kruskal() {
        int edgeSize = vertices.size() - 1;
        if (edgeSize == -1) return null;

        Set<EdgeInfo<V, E>> edgeInfos = new HashSet<>();
        // 小顶堆，选最小值 O(E)
        BinaryHeap<Edge<V, E>> heap = new BinaryHeap<>(edges, edgeComparator);
        // 并查集，判断是否形成环
        UnionFind<Vertex<V, E>> uf = new UnionFind<>();

        // 初始化并查集，每个顶点都是一个集合 O(V)
        vertices.forEach((V v, Vertex<V, E> vertex) -> {
            uf.makeSet(vertex);
        });

        // ElogE 次
        while (!heap.isEmpty() && edgeInfos.size() < edgeSize) {
            Edge<V, E> edge = heap.remove(); // O(logE)
            // 如果两个顶点已经联通，说明这条边会形成环，继续下一次循环 O(1)
            if (uf.isSame(edge.from, edge.to)) continue;

            edgeInfos.add(edge.info());
            uf.union(edge.from, edge.to);
        }

        return edgeInfos;
    }

    @Override
    public Map<V, PathInfo<V, E>> shortestPath(V begin) {
        return bellmanFord(begin);
    }

    private Map<V, PathInfo<V, E>> bellmanFord(V begin) {
        // 起点
        Vertex<V, E> beginVertex = vertices.get(begin);
        if (beginVertex == null) return null;
        // 存储最短路径的结果
        Map<V, PathInfo<V, E>> selectedPaths = new HashMap<>();
        // 标注起点和起点-起点的权值
        selectedPaths.put(begin, new PathInfo<>(weightManager.zero()));
        int count = vertices.size() - 1;
        for (int i = 0; i < count; i++) { // V - 1 次
            for (Edge<V, E> edge : edges) { // E 次
                // edge.from 已经是确定的最短路径了
                PathInfo<V, E> fromPath = selectedPaths.get(edge.from.value);
                if (fromPath == null) continue;
                relax(edge, fromPath, selectedPaths);
            }
        }

        // 第 V 次
        // 根据 bellmanFord 算法，第 V 次仍然可以松弛成功，说明有负权环
        for (Edge<V, E> edge : edges) {
            // edge.from 已经是确定的最短路径了
            PathInfo<V, E> fromPath = selectedPaths.get(edge.from.value);
            if (fromPath == null) continue;
            if (relax(edge, fromPath, selectedPaths)) {
                System.out.println("有负权环");
                return null;
            }
        }

        selectedPaths.remove(begin);
        return selectedPaths;
    }

    private boolean relax(Edge<V, E> edge, PathInfo<V, E> fromPath, Map<V, PathInfo<V, E>> selectedPaths) {
        // 新的可选择的最短路径：beginVertex 到 edge.from 的最短路径 + edge.weight
        E newWeight = weightManager.add(fromPath.getWeight(), edge.weight);
        // 以前的可选择的最短路径：beginVertex 到 edge.to 的最短路径
        PathInfo<V, E> oldPath = selectedPaths.get(edge.to.value);
        // 如果新的路径比以前的路径短，或者以前的路径为 null
        if (oldPath == null || weightManager.compare(newWeight, oldPath.getWeight()) > 0) {
            if (oldPath == null) {
                oldPath = new PathInfo<>();
                selectedPaths.put(edge.to.value, oldPath);
            } else {
                oldPath.getEdgeInfos().clear();
            }
            oldPath.setWeight(newWeight);
            oldPath.getEdgeInfos().addAll(fromPath.getEdgeInfos());
            oldPath.getEdgeInfos().add(edge.info());
            // 松弛成功
            return true;
        }
        // 松弛失败
        return false;
    }



    public Map<V, PathInfo<V, E>> dijkstra(V begin) {
        // 起点
        Vertex<V, E> beginVertex = vertices.get(begin);
        if (beginVertex == null) return null;
        // 存储最短路径的结果
        Map<V, PathInfo<V, E>> selectedPaths = new HashMap<>();
        // 存储还未确定是否是最短路径的边
        Map<Vertex<V, E>, PathInfo<V, E>> paths = new HashMap<>();
        // 初始化 paths
        beginVertex.outEdges.forEach((Edge<V, E> edge) -> {
            PathInfo<V, E> path = new PathInfo<>();
            path.setWeight(edge.weight);
            path.getEdgeInfos().add(edge.info());
            paths.put(edge.to, path);
        });
        while (!paths.isEmpty()) {
            // 从 paths 中挑一个最小的路径出来
            Map.Entry<Vertex<V, E>, PathInfo<V, E>> minEntry = getMinPath(paths);
            // minVertex 离开桌面
            Vertex<V, E> minVertex = minEntry.getKey();

            PathInfo<V, E> minPath = minEntry.getValue();
            // 将 minVertex 对应的最短路径放入 selectedPaths
            selectedPaths.put(minVertex.value, minPath);
            paths.remove(minVertex);

            // 对它的minVertex的outEdges进行松弛操作
            for (Edge<V, E> edge : minVertex.outEdges) {
                // 如果edge.to已经离开桌面，就没必要进行松弛操作
                if (selectedPaths.containsKey(edge.to.value)) continue;
                // 松弛操作
                relaxForDijkstra(edge, minPath, paths);
            }
        }
        selectedPaths.remove(begin);
        return selectedPaths;
    }


    // 从 paths 中挑一个最小的路径出来
    private Map.Entry<Vertex<V, E>, PathInfo<V, E>> getMinPath(Map<Vertex<V, E>, PathInfo<V, E>> paths) {
        Iterator<Map.Entry<Vertex<V, E>, PathInfo<V, E>>> it = paths.entrySet().iterator();
        Map.Entry<Vertex<V, E>, PathInfo<V, E>> minEntry = it.next();
        while (it.hasNext()) {
            Map.Entry<Vertex<V, E>, PathInfo<V, E>> entry = it.next();
            if (weightManager.compare(entry.getValue().getWeight(), minEntry.getValue().getWeight()) > 0) {
                minEntry = entry;
            }
        }
        return minEntry;
    }

    /**
     * 松弛操作
     *
     * @param edge     需要松弛的边
     * @param fromPath edge.from 刚刚被提起来的顶点的最短路径信息
     * @param paths    存放着还没有(离开桌面)的顶点的最短路径信息
     */
    private void relaxForDijkstra(Edge<V, E> edge, PathInfo<V, E> fromPath, Map<Vertex<V, E>, PathInfo<V, E>> paths) {
        // oldPath = beginVertex 到 edge.to 的最短路径信息
        PathInfo<V, E> oldPath = paths.get(edge.to);
        // 新的可选择的最短路径：beginVertex -> minVertex -> edge
        E newWeight = weightManager.add(fromPath.getWeight(), edge.weight);

        if (oldPath != null && weightManager.compare(newWeight, oldPath.getWeight()) <= 0) return;
        // 如果新的可选择的最短路径比旧的路径短，或者旧的路径为null
        if (oldPath == null) {
            oldPath = new PathInfo<>();
            paths.put(edge.to, oldPath);
        } else {
            oldPath.getEdgeInfos().clear();
        }
        oldPath.setWeight(newWeight);
        oldPath.getEdgeInfos().addAll(fromPath.getEdgeInfos());
        oldPath.getEdgeInfos().add(edge.info());
    }

    // Floyd算法
    @Override
    public Map<V, Map<V, PathInfo<V, E>>> shortestPath() {
        Map<V, Map<V, PathInfo<V, E>>> paths = new HashMap<>();
        // 初始化
        for (Edge<V, E> edge : edges) {
            // edge.from -> edge.to 的最短路径是 edge.weight
            Map<V, PathInfo<V, E>> map = paths.get(edge.from.value);
            if (map == null) {
                map = new HashMap<>();
                // paths 的 key 是起点， map 的 key 是终点
                paths.put(edge.from.value, map);
            }
            PathInfo<V, E> pathInfo = new PathInfo<>();
            pathInfo.setWeight(edge.weight);
            pathInfo.getEdgeInfos().add(edge.info());
            map.put(edge.to.value, pathInfo);
        }

        vertices.forEach((V v2, Vertex<V, E> vertext2) -> {
            vertices.forEach((V v1, Vertex<V, E> vertext1) -> {
                vertices.forEach((V v3, Vertex<V, E> vertext3) -> {
                    if (v1.equals(v2) || v2.equals(v3) || v1.equals(v3)) return;
                    // v1 -> v2
                    PathInfo<V, E> path12 = getPathInfo(v1, v2, paths);
                    if (path12 == null) return;
                    // v2 -> v3
                    PathInfo<V, E> path23 = getPathInfo(v2, v3, paths);
                    if (path23 == null) return;
                    // v1 -> v3
                    PathInfo<V, E> path13 = getPathInfo(v1, v3, paths);
                    E newWeight = weightManager.add(path12.getWeight(), path23.getWeight());
                    if (path13 != null && weightManager.compare(newWeight, path13.getWeight()) <= 0) return;
                    if (path13 == null) {
                        path13 = new PathInfo<>();
                        paths.get(v1).put(v3, path13);
                    } else {
                        path13.getEdgeInfos().clear();
                    }
                    path13.setWeight(newWeight);
                    path13.getEdgeInfos().addAll(path12.getEdgeInfos());
                    path13.getEdgeInfos().addAll(path23.getEdgeInfos());
                });
            });
        });

        return paths;
    }

    private PathInfo<V, E> getPathInfo(V from, V to, Map<V, Map<V, PathInfo<V, E>>> paths) {
        Map<V, PathInfo<V, E>> map = paths.get(from);
        return map == null ? null : map.get(to);
    }

    @Override
    public void print() {
        System.out.println("[顶点]-------------------");
        vertices.forEach((V v, Vertex<V, E> vertex) -> {
            System.out.println(v);
            System.out.println("out-----------");
            System.out.println(vertex.outEdges);
            System.out.println("in-----------");
            System.out.println(vertex.inEdges);
        });

        System.out.println("[边]-------------------");
        edges.forEach((Edge<V, E> edge) -> {
            System.out.println(edge);
        });
    }

    private static class Vertex<V, E> {
        V value;
        Set<Edge<V, E>> inEdges = new HashSet<>();
        Set<Edge<V, E>> outEdges = new HashSet<>();

        Vertex(V value) {
            this.value = value;
        }

        @Override
        public boolean equals(Object obj) {
            return Objects.equals(value, ((Vertex<V, E>) obj).value);
        }

        @Override
        public int hashCode() {
            return value == null ? 0 : value.hashCode();
        }

        @Override
        public String toString() {
            return value == null ? "null" : value.toString();
        }
    }

    private static class Edge<V, E> {
        Vertex<V, E> from;
        Vertex<V, E> to;
        E weight;

        Edge(Vertex<V, E> from, Vertex<V, E> to) {
            this.from = from;
            this.to = to;
        }

        @Override
        public boolean equals(Object obj) {
            Edge<V, E> edge = (Edge<V, E>) obj;
            return Objects.equals(from, edge.from) && Objects.equals(to, edge.to);
        }

        @Override
        public int hashCode() {
            return from.hashCode() * 31 + to.hashCode();
        }

        EdgeInfo<V, E> info() {
            return new EdgeInfo<>(from.value, to.value, weight);
        }

        @Override
        public String toString() {
            return from + "->" + to + ":" + weight;
        }
    }

}
